{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tools with Dependency Injection\n",
    "\n",
    "[Dependency Injection](https://en.wikipedia.org/wiki/Dependency_injection) is a secure way to connect external functions to agents without exposing sensitive data such as passwords, tokens, or personal information. This approach ensures that sensitive information remains protected while still allowing agents to perform their tasks effectively, even when working with large language models (LLMs).\n",
    "\n",
    "In this guide, we'll explore how to build secure workflows that handle sensitive data safely.\n",
    "\n",
    "As an example, we'll create an agent that retrieves user's account balance. The best part is that sensitive data like username and password are never shared with the LLM. Instead, it's securely injected directly into the function at runtime, keeping it safe while maintaining seamless functionality.\n",
    "\n",
    "\n",
    "## Why Dependency Injection Is Essential\n",
    "\n",
    "Here's why dependency injection is a game-changer for secure LLM workflows:\n",
    "\n",
    "- **Enhanced Security**: Your sensitive data is never directly exposed to the LLM.\n",
    "- **Simplified Development**: Secure data can be seamlessly accessed by functions without requiring complex configurations.\n",
    "- **Unmatched Flexibility**: It supports safe integration of diverse workflows, allowing you to scale and adapt with ease.\n",
    "\n",
    "In this guide, we'll explore how to set up dependency injection and build secure workflows. Let's dive in!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "To install `AG2`, simply run the following command:\n",
    "\n",
    "```bash\n",
    "pip install -U ag2[openai]\n",
    "```\n",
    "\n",
    "> **Note:** If you have been using `autogen` or `ag2`, all you need to do is upgrade it using:  \n",
    "> ```bash\n",
    "> pip install -U autogen\n",
    "> ```\n",
    "> or  \n",
    "> ```bash\n",
    "> pip install -U ag2\n",
    "> ```\n",
    "> as `autogen`, and `ag2` are aliases for the same PyPI package.  \n",
    "\n",
    "\n",
    "### Imports\n",
    "\n",
    "The functionality demonstrated in this guide is located in the `autogen.tools.dependency_injection` module. This module provides key components for dependency injection:\n",
    "\n",
    "- `BaseContext`: abstract base class used to define and encapsulate data contexts, such as user account information, which can then be injected into functions or agents securely.\n",
    "- `Depends`: a function used to declare and inject dependencies, either from a context (like `BaseContext`) or a function, ensuring sensitive data is provided securely without direct exposure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from typing import Annotated, Literal\n",
    "\n",
    "from pydantic import BaseModel\n",
    "\n",
    "from autogen import GroupChat, GroupChatManager\n",
    "from autogen.agentchat import ConversableAgent, UserProxyAgent\n",
    "from autogen.tools.dependency_injection import BaseContext, Depends"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define a BaseContext Class\n",
    "We start by defining a `BaseContext` class for accounts. This will act as the base structure for dependency injection. By using this approach, sensitive information like usernames and passwords is never exposed to the LLM."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Account(BaseContext, BaseModel):\n",
    "    username: str\n",
    "    password: str\n",
    "    currency: Literal[\"USD\", \"EUR\"] = \"USD\"\n",
    "\n",
    "\n",
    "alice_account = Account(username=\"alice\", password=\"password123\")\n",
    "bob_account = Account(username=\"bob\", password=\"password456\")\n",
    "\n",
    "account_ballace_dict = {\n",
    "    (alice_account.username, alice_account.password): 300,\n",
    "    (bob_account.username, bob_account.password): 200,\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Helper Functions\n",
    "To ensure that the provided account is valid and retrieve its balance, we create two helper functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def _verify_account(account: Account):\n",
    "    if (account.username, account.password) not in account_ballace_dict:\n",
    "        raise ValueError(\"Invalid username or password\")\n",
    "\n",
    "\n",
    "def _get_balance(account: Account):\n",
    "    _verify_account(account)\n",
    "    return f\"Your balance is {account_ballace_dict[(account.username, account.password)]}{account.currency}\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Agent Configuration\n",
    "\n",
    "Configure the agents for the interaction.\n",
    "\n",
    "- `config_list` defines the LLM configurations, including the model and API key.\n",
    "- `UserProxyAgent` simulates user inputs without requiring actual human interaction (set to `NEVER`).\n",
    "- `AssistantAgent` represents the AI agent, configured with the LLM settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = [{\"api_type\": \"openai\", \"model\": \"gpt-5-nano\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}]\n",
    "\n",
    "assistant = ConversableAgent(\n",
    "    name=\"assistant\",\n",
    "    llm_config={\"config_list\": config_list},\n",
    ")\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy_1\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    llm_config=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Injecting a BaseContext Parameter\n",
    "\n",
    "In the example below we register the function and use dependency injection to automatically inject the bob_account Account object into the function. This `account` parameter will not be visible to the LLM.\n",
    "\n",
    "**Note:** You can also use `account: Account = Depends(bob_account)` as an alternative syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@user_proxy.register_for_execution()\n",
    "@assistant.register_for_llm(description=\"Get the balance of the account\")\n",
    "def get_balance_1(\n",
    "    # Account which will be injected to the function\n",
    "    account: Annotated[Account, Depends(bob_account)],\n",
    "    # It is also possible to use the following syntax to define the dependency\n",
    "    # account: Account = Depends(bob_account),\n",
    ") -> str:\n",
    "    return _get_balance(account)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we initiate a chat to retrieve the balance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_proxy.initiate_chat(assistant, message=\"Get the user's account balance\", max_turns=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Injecting Parameters Without BaseContext\n",
    "\n",
    "Sometimes, you might not want to use `BaseContext`. Here's how to inject simple parameters directly.\n",
    "\n",
    "### Agent Configuration\n",
    "\n",
    "Configure the agents for the interaction.\n",
    "\n",
    "- `config_list` defines the LLM configurations, including the model and API key.\n",
    "- `UserProxyAgent` simulates user inputs without requiring actual human interaction (set to `NEVER`).\n",
    "- `AssistantAgent` represents the AI agent, configured with the LLM settings."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = [{\"api_type\": \"openai\", \"model\": \"gpt-5-nano\", \"api_key\": os.environ[\"OPENAI_API_KEY\"]}]\n",
    "assistant = ConversableAgent(\n",
    "    name=\"assistant\",\n",
    "    llm_config={\"config_list\": config_list},\n",
    ")\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy_1\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    llm_config=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Register the Function with Direct Parameter Injection\n",
    "Instead of injecting a full context like `Account`, you can directly inject individual parameters, such as the username and password, into a function. This allows for more granular control over the data injected into the function, and still ensures that sensitive information is managed securely.\n",
    "\n",
    "Here’s how you can set it up:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_username() -> str:\n",
    "    return \"bob\"\n",
    "\n",
    "\n",
    "def get_password() -> str:\n",
    "    return \"password456\"\n",
    "\n",
    "\n",
    "@user_proxy.register_for_execution()\n",
    "@assistant.register_for_llm(description=\"Get the balance of the account\")\n",
    "def get_balance_2(\n",
    "    username: Annotated[str, Depends(get_username)],\n",
    "    password: Annotated[str, Depends(get_password)],\n",
    "    # or use lambdas\n",
    "    # username: Annotated[str, Depends(lambda: \"bob\")],\n",
    "    # password: Annotated[str, Depends(lambda: \"password456\")],\n",
    ") -> str:\n",
    "    account = Account(username=username, password=password)\n",
    "    return _get_balance(account)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initiate the Chat\n",
    "As before, initiate a chat to test the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "user_proxy.initiate_chat(assistant, message=\"Get users balance\", max_turns=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Aligning Contexts to Agents\n",
    "\n",
    "You can match specific dependencies, such as 3rd party system credentials, with specific agents by using tools with dependency injection.\n",
    "\n",
    "In this example we have 2 external systems and have 2 related login credentials. We don't want or need the LLM to be aware of these credentials.\n",
    "\n",
    "### Mock third party systems"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mock third party system functions\n",
    "# Imagine that these use the username and password to authenticate\n",
    "\n",
    "\n",
    "def weather_api_call(username: str, password: str, location: str) -> str:\n",
    "    print(f\"Accessing third party Weather System using username {username}\")\n",
    "    return \"It's sunny and 40 degrees Celsius in Sydney, Australia.\"\n",
    "\n",
    "\n",
    "def my_ticketing_system_availability(username: str, password: str, concert: str) -> bool:\n",
    "    print(f\"Accessing third party Ticketing System using username {username}\")\n",
    "    return False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weather_agent = ConversableAgent(\n",
    "    name=\"weather_agent\",\n",
    "    system_message=\"You are a Weather Agent, you can only get the weather.\",\n",
    "    description=\"Weather Agent solely used for getting weather.\",\n",
    "    llm_config={\"config_list\": config_list},\n",
    ")\n",
    "\n",
    "ticket_agent = ConversableAgent(\n",
    "    name=\"ticket_agent\",\n",
    "    system_message=\"You are a Ticketing Agent, you can only get ticket availability.\",\n",
    "    description=\"Ticketing Agent solely used for getting ticket availability.\",\n",
    "    llm_config={\"config_list\": config_list},\n",
    ")\n",
    "user_proxy = UserProxyAgent(\n",
    "    name=\"user_proxy_1\",\n",
    "    human_input_mode=\"NEVER\",\n",
    "    llm_config=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create BaseContext class for credentials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We create a class based on BaseContext\n",
    "class ThirdPartyCredentials(BaseContext, BaseModel):\n",
    "    username: str\n",
    "    password: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Credentials and Functions with Dependency Injection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Weather API\n",
    "weather_account = ThirdPartyCredentials(username=\"ag2weather\", password=\"wbkvEehV1A\")\n",
    "\n",
    "\n",
    "@user_proxy.register_for_execution()\n",
    "@weather_agent.register_for_llm(description=\"Get the weather for a location\")\n",
    "def get_weather(\n",
    "    location: str,\n",
    "    credentials: Annotated[ThirdPartyCredentials, Depends(weather_account)],\n",
    ") -> str:\n",
    "    # Access the Weather API using the credentials\n",
    "    return weather_api_call(username=credentials.username, password=credentials.password, location=location)\n",
    "\n",
    "\n",
    "# Ticketing System\n",
    "ticket_system_account = ThirdPartyCredentials(username=\"ag2tickets\", password=\"EZRIVeVWvA\")\n",
    "\n",
    "\n",
    "@user_proxy.register_for_execution()\n",
    "@ticket_agent.register_for_llm(description=\"Get the availability of tickets for a concert\")\n",
    "def tickets_available(\n",
    "    concert_name: str,\n",
    "    credentials: Annotated[ThirdPartyCredentials, Depends(ticket_system_account)],\n",
    ") -> bool:\n",
    "    return my_ticketing_system_availability(\n",
    "        username=credentials.username, password=credentials.password, concert=concert_name\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Group Chat and run"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "groupchat = GroupChat(agents=[user_proxy, weather_agent, ticket_agent], messages=[], max_round=5)\n",
    "manager = GroupChatManager(groupchat=groupchat, llm_config={\"config_list\": config_list})\n",
    "\n",
    "message = (\n",
    "    \"Start by getting the weather for Sydney, Australia, and follow that up by checking \"\n",
    "    \"if there are tickets for the 'AG2 Live' concert.\"\n",
    ")\n",
    "user_proxy.initiate_chat(manager, message=message, max_turns=1)"
   ]
  }
 ],
 "metadata": {
  "front_matter": {
   "description": "Tools Dependency Injection",
   "tags": [
    "tools",
    "dependency injection",
    "function calling"
   ]
  },
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
